# Visibility Analysis: Visualization Functions
# Alex F. Gazmararian  
# agazmararian@gmail.com

#' Create state-level map with survey points and projects
#' @param state.boundary sf object with state boundaries
#' @param survey.points sf object with survey respondent locations
#' @param eia.projects sf object with EIA projects (optional)
#' @param bgm.projects sf object with BGM projects (optional)
#' @param title Plot title (optional)
#' @param base.size Base font size (default: 12)
#' @param point.size Point size for survey locations (default: 1)
#' @return ggplot object
create_state_map <- function(state.boundary, survey.points, eia.projects = NULL, 
                           bgm.projects = NULL, title = NULL, base.size = 12, point.size = 1) {
  p <- ggplot2::ggplot() +
    ggplot2::geom_sf(data = state.boundary, fill = "white", color = "black", size = 0.3) +
    ggplot2::geom_sf(data = survey.points, color = "blue", size = point.size, alpha = 0.6) +
    ggplot2::theme_void(base_size = base.size) +
    ggplot2::theme(
      panel.grid = ggplot2::element_blank(),
      axis.text = ggplot2::element_blank(),
      axis.ticks = ggplot2::element_blank()
    )
  
  # Add EIA projects if provided
  if (!is.null(eia.projects) && nrow(eia.projects) > 0) {
    p <- p + ggplot2::geom_sf(data = eia.projects, color = "red", size = 0.8, alpha = 0.7)
  }
  
  # Add BGM projects if provided
  if (!is.null(bgm.projects) && nrow(bgm.projects) > 0) {
    p <- p + ggplot2::geom_sf(data = bgm.projects, color = "green", size = 0.8, alpha = 0.7)
  }
  
  # Add title if provided
  if (!is.null(title)) {
    p <- p + ggplot2::labs(title = title)
  }
  
  return(p)
}

#' Create distance plot showing buffer zones around survey point
#' @param state.boundary sf object with state boundaries
#' @param survey.point sf object with single survey point
#' @param eia.projects sf object with EIA projects
#' @param bgm.projects sf object with BGM projects
#' @param buffer.distances Numeric vector of buffer distances in km (default: c(15, 25, 50, 100))
#' @param coord.labels Character labels for coordinates (optional)
#' @param filename Output filename for saving (optional)
#' @return ggplot object
create_distance_plot <- function(state.boundary, survey.point, eia.projects, bgm.projects,
                               buffer.distances = c(15, 25, 50, 100), coord.labels = NULL,
                               filename = NULL) {
  
  # Create buffer zones
  buffers <- purrr::map(buffer.distances, ~ {
    sf::st_buffer(survey.point, dist = .x * 1000) # Convert km to meters
  }) %>%
    purrr::set_names(paste0(buffer.distances, "km"))
  
  # Start with base map
  p <- ggplot2::ggplot() +
    ggplot2::geom_sf(data = state.boundary, fill = "white", color = "black", size = 0.3)
  
  # Add buffer zones (largest first so they layer properly)
  for (i in length(buffers):1) {
    buffer_name <- names(buffers)[i]
    alpha_val <- 0.1 + (i - 1) * 0.05  # Increasing alpha for smaller buffers
    p <- p + ggplot2::geom_sf(data = buffers[[i]], fill = "blue", alpha = alpha_val, color = "blue", size = 0.3)
  }
  
  # Add projects
  if (nrow(eia.projects) > 0) {
    p <- p + ggplot2::geom_sf(data = eia.projects, color = "red", size = 1.5, alpha = 0.8)
  }
  if (nrow(bgm.projects) > 0) {
    p <- p + ggplot2::geom_sf(data = bgm.projects, color = "green", size = 1.5, alpha = 0.8)
  }
  
  # Add survey point (on top)
  p <- p + ggplot2::geom_sf(data = survey.point, color = "black", size = 3, shape = 21, fill = "yellow")
  
  # Style the plot
  p <- p +
    ggplot2::theme_void() +
    ggplot2::theme(
      panel.grid = ggplot2::element_blank(),
      axis.text = ggplot2::element_blank(),
      axis.ticks = ggplot2::element_blank(),
      legend.position = "bottom"
    )
  
  # Add coordinate labels if provided
  if (!is.null(coord.labels)) {
    coords <- get_coords(survey.point)
    p <- p + ggplot2::labs(
      subtitle = paste0("Location: ", coord.labels, " (", round(coords$lon, 3), ", ", round(coords$lat, 3), ")")
    )
  }
  
  # Save if filename provided  
  if (!is.null(filename)) {
    ggplot2::ggsave(filename, plot = p, width = 10, height = 8, dpi = 300)
  }
  
  return(p)
}

#' Extract coordinates from sf object
#' @param input sf object
#' @return List with lon, lat, and bounding box info
get_coords <- function(input) {
  coords <- sf::st_coordinates(input)
  min_y <- min(coords[, 2])
  max_y <- max(coords[, 2])
  min_x <- min(coords[, 1])
  max_x <- max(coords[, 1])
  
  list(
    lon = mean(c(min_x, max_x)),
    lat = mean(c(min_y, max_y)),
    bbox = list(
      xmin = min_x, xmax = max_x,
      ymin = min_y, ymax = max_y
    )
  )
}

#' Plot heterogeneous treatment effects
#' @param ame.model List of marginal effects models
#' @param hetvar Variable name for heterogeneity analysis
#' @param modvarlab Label for the moderating variable
#' @return ggplot object
plot_het_effects <- function(ame.model, hetvar, modvarlab) {
  p.het <- ame.model %>%
    dplyr::bind_rows(.id = "model") %>%
    dplyr::filter(term == hetvar) %>%
    ggplot2::ggplot(ggplot2::aes(x = value, y = estimate)) +
    ggplot2::geom_point(size = 2) +
    ggplot2::geom_errorbar(ggplot2::aes(ymin = conf.low, ymax = conf.high), width = 0.1) +
    ggplot2::geom_hline(yintercept = 0, linetype = "dashed", color = "red", alpha = 0.7) +
    ggplot2::labs(
      x = modvarlab,
      y = "Marginal Effect",
      title = paste("Heterogeneous Effects by", modvarlab)
    ) +
    ggplot2::theme_minimal() +
    ggplot2::theme(
      panel.grid.minor = ggplot2::element_blank()
    )
  
  return(p.het)
}
